import os
import time

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from model import TTSR
from loss.loss import get_loss_dict
from torch.utils.data import Dataset, DataLoader

from model import LTE
# from utils import mkExpDir
from trainer import Trainer
from dataloader.dataset import ImageDataset

import warnings
warnings.filterwarnings('ignore')


if __name__ == '__main__':
    cpu = torch.device('cpu')
    device = torch.device('cuda')
    _model = TTSR.TTSR().to(device)
    t = Trainer(_model)

    ## load data
    data = torch.load('/home/xincz/Documents/Tensors/tensor.pt')
    tar_ds = ImageDataset(list(data.keys()), data)
    tar_dl = DataLoader(tar_ds, shuffle=False, batch_size=64, num_workers=4)

    # Preprocess the targets
    print('start preprocessing ==========')
    tar_names = list(data.keys())
    since = time.time()
    targets = torch.cat([F.normalize(  # targets: [9996, 176, 1152]
        F.unfold(tar.to(device), kernel_size=(3, 3), padding=1, stride=3) \
            .permute(0, 2, 1), dim=2).to(cpu) for tar, _ in tar_dl], dim=0)
    print("targets preprocessing elapsed: ", time.time() - since)

    del data
    del tar_ds
    del tar_dl

    # torch.save({
    #     "tar_names": tar_names,
    #     "targets": targets
    # }, "/home/xincz/Documents/Tensors/targets-info.pt")
    #
    # data = torch.load("/home/xincz/Documents/Tensors/targets-info.pt")
    # tar_names = data["tar_names"]
    # targets = data["targets"]
    # print(targets.shape)

    # begin search test
    indir = '/home/xincz/Documents/Datasets/models_540'
    s_time = time.time()
    res = t.search_list(indir, 'models_540-1', {
        "tar_names": tar_names,
        "targets": targets
    })

    # ## 一对比较
    # s_time = time.time()
    # ref = '/home/hjj/Downloads/test/_122/465139.jpg'
    # lr_path =  '/home/hjj/Downloads/test/yy搜索测试/其他/todo/465207.png'
    # print(ref,'--------')
    # # k = 4
    # res = t.test_pair(ref,lr_path, 2)
    # print(time.time() - s_time)

    # # # 文件夹search
    # # for file_name in ['波点','卡通','抽象','民族风','几何','其他']:
    # for file_name in ['卡通']:
    #     print('===========:',file_name)
    #     indir = '/home/hjj/Downloads/test/yy测试花型/'+file_name + '/todo'
    #     s_time = time.time()
    #     res = t.search_list(indir,file_name,'ref-4-4-conv10.npy')
    #     print(time.time() - s_time)

    # ## 单张search
    # s_time = time.time()
    # input_path = '/home/hjj/Downloads/test/yy搜索测试/其他/todo/465233.png'
    # print(input_path,'--------')
    # res = t.search(input_path,'ref-4-4-conv10.npy')
    # print(time.time() - s_time)

    # ## 单张ref
    # s_time = time.time()
    # input_path = '/home/hjj/Downloads/test/_122/465207.jpg'
    # print(input_path,'--------')
    # # k = 4
    # res = t.ref_one(input_path, 6)
    # print(time.time() - s_time)

    # indir = '_122'
    # inlist1 = os.listdir(indir)
    # inlist1.sort()
    # inlist1 = inlist1[:4600]
    # inlist1 = [os.path.join(indir,i) for i in inlist1]
    #
    # indir = '小午花型'
    # inlist2 = os.listdir(indir)
    # inlist2.sort()
    # inlist2 = [os.path.join(indir,i) for i in inlist2]
    #
    # inlist = inlist1+inlist2
    # k = 4 # 4: 16倍, 3: 12倍
    # res = t.ref(inlist,k)
    # np.save('model_122.npy', res)
